# enhance_video.py

import os
import time
import glob
import shutil

# Dynamically determine the project root directory
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

def check_and_download_models():
    """Checks if the models are present and downloads them if not."""
    realesrgan_model_path = os.path.join(PROJECT_ROOT, 'Real-ESRGAN', 'weights', 'RealESRGAN_x4plus.pth')
    rife_model_path = os.path.join(PROJECT_ROOT, 'RIFE', 'train_log', 'flownet.pkl')
    if not (os.path.exists(realesrgan_model_path) and os.path.exists(rife_model_path)):
        print("Models not found. Starting download...")
        # Ensure download script is called with the correct context
        subprocess.run(['python3', os.path.join(PROJECT_ROOT, 'download_models.py')])

def enhance_pipeline(input_video_path):
    """
    Receives a video path, performs upscaling and frame interpolation, and returns the final file path.
    """
    # Gracefully handle the initial call from Gradio with None input
    if not input_video_path:
        print("No input video provided. Waiting for upload...")
        return None

    check_and_download_models()
    
    print(f"开始处理视频: {input_video_path}")

    # --- 0. Set up environment variables and paths ---
    python_path_env = (
        f"PYTHONPATH={os.path.join(PROJECT_ROOT, 'BasicSR')}:"
        f"{os.path.join(PROJECT_ROOT, 'Real-ESRGAN')}:"
        f"{os.path.join(PROJECT_ROOT, 'RIFE')}:"
        f"$PYTHONPATH"
    )

    # --- 1. Set up temporary and output file paths ---
    timestamp = int(time.time())
    tmp_dir = "/tmp/video-improvement"
    os.makedirs(tmp_dir, exist_ok=True)
    
    realesrgan_output_dir = os.path.join(tmp_dir, f"realesrgan_out_{timestamp}")
    os.makedirs(realesrgan_output_dir, exist_ok=True)

    upscaled_video_path = os.path.join(tmp_dir, f"upscaled_{timestamp}.mp4")
    output_video_path = os.path.join(PROJECT_ROOT, 'output', f"enhanced_{timestamp}.mp4")

    os.makedirs(os.path.join(PROJECT_ROOT, 'output'), exist_ok=True)

    # --- 2. Execute Real-ESRGAN script ---
    if os.environ.get('SKIP_SUPER_RESOLUTION'):
        print("步骤 1/2: 正在进行视频超分辨率 (已通过环境变量禁用)...")
        upscaled_video_path = input_video_path
        print("超分已禁用，直接使用原始视频进行补帧。")
    else:
        print("步骤 1/2: 正在进行视频超分辨率...")
        esrgan_script_path = os.path.join(PROJECT_ROOT, 'Real-ESRGAN', 'inference_realesrgan_video.py')
        esrgan_command = (
            f'{python_path_env} python "{esrgan_script_path}" '
            f'-i "{input_video_path}" -o "{realesrgan_output_dir}" '
            f'--suffix outx4 --face_enhance'
        )
        os.system(esrgan_command)
        print("超分完成！")

        # --- 2.1. Find and rename the Real-ESRGAN output file ---
        try:
            generated_files = glob.glob(os.path.join(realesrgan_output_dir, '*.mp4'))
            if generated_files:
                shutil.move(generated_files[0], upscaled_video_path)
                shutil.rmtree(realesrgan_output_dir)
            else:
                raise FileNotFoundError("Real-ESRGAN did not produce an output file.")
        except Exception as e:
            print(f"Error handling Real-ESRGAN output: {e}")
            return None
    # --- 3. Execute RIFE script ---
    if not os.path.exists(upscaled_video_path):
        print("Upscaled video not found, skipping RIFE.")
        return None
        
    print("步骤 2/2: 正在进行视频补帧...")
    rife_dir = os.path.join(PROJECT_ROOT, 'RIFE')
    rife_script_path = os.path.join(rife_dir, 'inference_video.py')
    rife_command = (
        f'cd "{rife_dir}" && '
        f'{python_path_env} python "{rife_script_path}" --video "{upscaled_video_path}" --output "{output_video_path}"'
    )
    os.system(rife_command)
    print("补帧完成！")

    # --- 4. Clean up temporary files ---
    if os.path.exists(upscaled_video_path):
        os.remove(upscaled_video_path)

    print(f"处理完毕，文件已保存至: {output_video_path}")
    return output_video_path
